import sys
import json
from IPython.display import Markdown, display
from llama_index.llms.ollama import Ollama
from llama_index.core.retrievers import VectorIndexRetriever
from llama_index.core.query_engine import RetrieverQueryEngine
from llama_index.core.postprocessor import SimilarityPostprocessor
from llama_index.embeddings.fastembed import FastEmbedEmbedding
from llama_index.core import SQLDatabase
from llama_index.core import Settings
from llama_index.core.query_engine import NLSQLTableQueryEngine
from sqlalchemy import create_engine
import openai
import logging

# from actions.api.mysql_tab import MysqlTab
from mysql_tab import MysqlTab
logger = logging.getLogger(__name__)

class RagMysql:
    """ RagMysql """

    def __init__(self ):
        logger.info(' RagMysql -init')
        
        openai.api_key = '？？？'
        openai.base_url = '？？？'
        from llama_index.llms.openai import OpenAI
        self.llm = OpenAI(temperature=0.1, model="gpt-3.5-turbo")
        # self.llm = Ollama(model="codeqwen:7b-chat", base_url="http://localhost:11434",request_timeout=60000.0)
        embed_model = FastEmbedEmbedding(model_name="BAAI/bge-small-zh-v1.5")        
        Settings.llm = self.llm
        Settings.embed_model = embed_model
        engine = create_engine('mysql+mysqlconnector://testfrom:testfrom@8.130.66.117:3306/testfrom')
        # sql_database = SQLDatabase(engine, include_tables=["system_users","system_user_role"])
        self.sql_database = SQLDatabase(engine)
        
    def rag(self,query_str):
        prompt=f'''
        提取下面句子的名词。结果采用逗号分隔。
        句子是:{query_str}
        请只回复答案，不要写解释。
        '''
        resp = self.llm.complete(prompt)
        logger.info('#########################')
        logger.info(resp)
        
        tabs = MysqlTab()
        rel_tabs=tabs.getTabs(resp.text)
        logger.info(rel_tabs)
        tabs.close()
        
        nlp_engine = NLSQLTableQueryEngine(
            sql_database=self.sql_database, 
            tables=rel_tabs, 
            llm=self.llm,
            verbose=False
        )
        logger.info('#########################')
        response = nlp_engine.query(query_str)
        logger.info(f"SQL query: {response.metadata['sql_query']}")
        return response.metadata['sql_query']
        
    def readDoc(self):
        from llama_index.core import VectorStoreIndex
        from llama_index.vector_stores.chroma import ChromaVectorStore
        import chromadb
        db2 = chromadb.PersistentClient(path="/home/softrobot/kb/chromadb")
        chroma_collection = db2.get_or_create_collection("docs")
        vector_store = ChromaVectorStore(chroma_collection=chroma_collection)
        index = VectorStoreIndex.from_vector_store(
            vector_store,
            embed_model=self.embed_model,
        )

        # query
        prompt =f'''
        {query_str}
        找出相似的问题和sql语句
        '''
        doc_engine = index.as_query_engine(similarity_top_k=2)
        rag_doc = doc_engine.query(prompt)
        print('#########################')
        print(rag_doc)
        from llama_index.core import  PromptTemplate
        from llama_index.core.prompts import PromptType
        TEXT_TO_SQL_TMPL = f'''
            Given an input question, first create a syntactically correct SQL 
            query to run, then look at the results of the query and return the answer. 
            You can order the results by a relevant column to return the most 
            interesting examples in the database.\n\n
            Never query for all the columns from a specific table, only ask for a 
            few relevant columns given the question.\n\n
            Pay attention to use only the column names that you can see in the schema 
            description. 
            Be careful to not query for columns that do not exist. 
            Pay attention to which column is in which table. 
            Also, qualify column names with the table name when needed. \n
            If needing to group on Array Columns use the ClickHouse function arrayJoin e.g. arrayJoin(columnName) \n
            For example, the following query identifies the most popular database:\n
            {rag_doc}     
            You are required to use the following format, each taking one line:\n\n
            Question: Question here\n
            SQLQuery: SQL Query to run\n
            SQLResult: Result of the SQLQuery\n
            Only use tables listed below.\n
            Question: {query_str}\n
            SQLQuery:
        '''
        TEXT_TO_SQL_PROMPT = PromptTemplate(
            TEXT_TO_SQL_TMPL,
            prompt_type=PromptType.TEXT_TO_SQL,
        )


ragMysql = RagMysql()
# ret = ragMysql.rag("角色id大于1,并且用户名等于系统管理员，这个用户的备注是什么")
ret = ragMysql.rag("1+1")
print(ret)